import time
import json
import glob
import struct
import zipfile
import numpy as np
import xarray as xr
import pandas as pd
from osgeo import gdal
from typing import Any
from typing import Dict
from typing import Union
from pathlib import Path
from pycwr.io import read_auto
from utils.logger import Logger

logger = Logger(filename=f'./log/xgrid_writer.log', level='info', when='D', back_count=1).logger

_type_mapping_to_xgrid = {
    'float32': ('Float32', 999999.0),
    'uint8': ('Int8', 255),
    'uint16': ('Int16', 9999),
}


def xgrid_writer(data: Union[xr.DataArray, xr.Dataset],
                 timestamp,
                 save_path: Union[str, Path],
                 dtype: str = 'uint8',
                 dataOffset: float = None,
                 dataScale: float = None,
                 units: str = None,
                 undef: float = None,
                 meta: Dict[str, Any] = None,
                 ) -> Path:
    """将格点数据打包成指定zip压缩的指定二进制格式"""

    timeList = data.time.values.tolist()
    level_list = data.level.values.tolist()
    lon = data.lon.values.tolist()
    lat = data.lat.values.tolist()

    file_name = Path(save_path)
    logger.info(f'写入zip [{str(file_name)}]')

    if not file_name.parent.exists():
        file_name.parent.mkdir(parents=True)

    data_type, default_undef = _type_mapping_to_xgrid[dtype]
    if undef is None:
        undef = default_undef

    with zipfile.ZipFile(file_name, mode='w', compression=zipfile.ZIP_DEFLATED, compresslevel=9) as zfile:
        dfile_name = '.'.join(file_name.name.split('.')[:-1]) + '.dat'
        with zfile.open(dfile_name, mode='w') as raw:
            header = {
                'xSize': len(lon),
                'xStart': lon[0],
                'xEnd': lon[-1],
                'xDelta': round(lon[1] - lon[0], 2),

                'ySize': len(lat),
                'yStart': lat[0],
                'yEnd': lat[-1],
                'yDelta': round(lat[1] - lat[0], 2),

                'levels': len(level_list),
                'levelList': level_list,

                'times': len(timeList),
                'timeList': timeList,
                'timestamp': timestamp,

                'undef': undef,
                'dataScale': dataScale,
                'dataOffset': dataOffset,
                'units': units,
                'littleEndian': True,
                'dataType': data_type,
                'unsigned': dtype.startswith('u')
            }
            if meta is not None:
                header.update(meta)
            target_array: np.ndarray = data.values

            if dataOffset is not None:
                target_array -= dataOffset

            if dataScale is not None:
                target_array /= dataScale

            target_array = np.nan_to_num(target_array, nan=undef)

            header = json.dumps(header, ensure_ascii=False).encode('utf-8')
            header_len = len(header)
            raw.write(struct.pack('i', header_len))
            raw.write(header)
            raw.write(target_array.astype(dtype).tobytes())
    logger.info('zip写入完成')
    return file_name


def get_radar(rfile, stride_x=320, stride_y=320):
    level = np.linspace(2000, 14000, 7).astype(int).tolist()

    sa = read_auto(rfile)
    rfile = Path(rfile)
    end_time = sa.scan_info.end_time.values
    timestamp = pd.to_datetime(end_time).timestamp().real
    times = [end_time]

    lat, lon = int(sa.scan_info['latitude'] * 100), int(sa.scan_info['longitude'] * 100)

    lon1d = np.arange(lon - stride_x, lon + stride_x) / 100
    lat1d = np.arange(lat + stride_y, lat - stride_y, step=-1) / 100

    rdata = np.zeros([len(times), len(level), len(lon1d), len(lat1d)])
    for i, height in enumerate(level):
        sa.add_product_CAPPI_lonlat(XLon=lon1d, YLat=lat1d, level_height=height)
        k = list(sa.product.keys())[-1]
        rdata[0][i] = sa.product[k].values

    rdar_xarr = xr.DataArray(rdata, coords=[times, level, lon1d, lat1d], dims=['time', 'level', 'lon', 'lat'])

    save_path = Path('./results/', rfile.parent.name, rfile.name.split('.')[0] + '.zip')

    xgrid_writer(rdar_xarr, timestamp=timestamp, save_path=save_path, dtype='uint8', dataOffset=-50, dataScale=0.5,
                 units='dBZ')


def get_UV_PRE(tfile):
    ds = gdal.Open(tfile)

    tfile = Path(tfile)
    ftime = '20' + tfile.name[:8]
    times = [ftime]
    row = ds.RasterXSize
    col = ds.RasterYSize
    band = ds.RasterCount
    lon0, lon_delta, _lon, lat0, _lat, lat_delta = ds.GetGeoTransform()

    lon1d = [round(lon0 + lon_delta * i, 3) for i in range(row)]
    lat1d = [round(lat0 + lat_delta * i, 3) for i in range(col)]

    if band == 1:
        tif_data = ds.ReadAsArray()
        data = tif_data[np.newaxis, np.newaxis]
    else:
        data = np.zeros([1, 1, row, col])
        for i in range(band):
            dt = ds.GetRasterBand(1)
            data[0, i, :, :] = dt.ReadAsArray(0, 0, row, col)

    level = [i + 1 for i in range(band)]
    data /= 10
    dataArr = xr.DataArray(data, coords=[times, level, lat1d, lon1d], dims=['time', 'level', 'lat', 'lon'])

    return dataArr


if __name__ == '__main__':
    # fi = '/home/gym/projects/nowcast_shaoxing/data/tif_PRE/21072400_wgs84.tif'

    xfiles = glob.glob('/home/gym/projects/nowcast_shaoxing/data/tif_MaxWind*/210724*.tif')
    xfiles.sort()
    data0 = 0
    for xf in xfiles:
        xdata = get_UV_PRE(xf)

        tfile = Path(xf)

        ftime = '20' + tfile.name[:8]
        timestp = time.mktime(time.strptime(ftime, '%Y%m%d%H'))

        data_offset = None
        unit = None

        if 'PRE' in tfile.parent.name:
            xdata.values += data0
            data0 = xdata.copy()
            unit = 'mm'

        elif 'Wind' in tfile.parent.name:
            data_offset = -50.
            unit = 'm/s'

        savePath = Path('./results/', tfile.parent.name, tfile.name.split('.')[0] + '.zip')

        xgrid_writer(xdata, timestamp=timestp, save_path=savePath, dtype='uint16', dataOffset=data_offset,
                     dataScale=0.1, units=unit)

    # rdar_file = r"/home/spd2/datasets/gym/datasets/shaoxing/radar/2020-08-03/Z_RADR_I_Z9574_20200803212400_O_DOR_SA_CAP.bin.bz2"

    # xfiles = glob.glob('home/spd2/datasets/gym/datasets/shaoxing/radar/2020-08-03/*.bz2')
    # get_radar(rdar_file, stride_x=320, stride_y=320)
